import pandas as pd
import numpy as np
from pandasai import SmartDataframe
from pandasai.llm import OpenAI
import os
# 初始化 LLM
os.environ["REQUESTS_CA_BUNDLE"] = ""
os.environ["SSL_CERT_FILE"] = ""


class MYOPENAI(OpenAI):
    

    _supported_chat_models = [
        "gpt-4",
        "gpt-4-0613",
        "gpt-4-32k",
        "gpt-4-32k-0613",
        "gpt-4-1106-preview",
        "gpt-3.5-turbo",
        "gpt-3.5-turbo-16k",
        "gpt-3.5-turbo-0613",
        "gpt-3.5-turbo-16k-0613",
        "gpt-3.5-turbo-1106",
        'Qwen/QwQ-32B'
        
    ]


# api_token=os.environ["OPENAI_API_KEY"]
#MyLLM = OpenAI(api_token=api_token,openai_proxy='')  
MyLLM = MYOPENAI(api_token="sk-lufpkjmnhkzdfuuosshszezhnjldzyvwvjseflxbatvtzmdl", api_base="https://api.siliconflow.cn/v1",model='Qwen/QwQ-32B')  


# 随机生成一些员工数据
np.random.seed(42)  # 设置随机种子以保证结果可复现

# 员工信息
data = {
    "员工ID": np.arange(1, 11),  # 员工编号
    "姓名": ["张三", "李四", "王五", "赵六", "孙七", "周八", "吴九", "郑十", "钱伯", "孔仲"],
    "性别": np.random.choice(["男", "女"], size=10),  # 随机分配性别
    "年龄": np.random.randint(20, 60, size=10),  # 随机生成年龄
    "职位": np.random.choice(["护林员", "导游", "管理员", "清洁工"], size=10),  # 随机分配职位
    "入职日期": pd.date_range(start="2015-01-01", periods=10, freq="M"),  # 随机生成入职日期
    "工资": np.random.randint(3000, 8000, size=10)  # 随机生成工资
}

# 创建 DataFrame
df = pd.DataFrame(data)

# 显示数据
print("国家森林公园员工信息表：")
print(df)

# 初始化 SmartDataframe
smart_df = SmartDataframe(df, config={
    "llm": MyLLM,
    "verbose": True,
    "security": None,
    "enforce_privacy": False,
    "enable_cache": False,
    "custom_whitelisted_dependencies": ["scikit-learn", 'plotly', 'matplotlib', 'numpy', 'pandas', 'io', 'sys', 'chr', 'glob', 'b64decoder', 'collections'],
    "save_charts": True,
    "data_viz_library": "plotly"
})

# 定义预设问题
preset_questions = {
    1: "用饼图分析员工职位分布",
    2: "显示员工年龄的分布情况",
    3: "计算每个职位的平均工资",
    4: "显示入职日期最早的员工信息",
    5: "计算员工工资的总和和平均值",
    6: "根据性别统计员工数量"
}

# 显示预设问题供用户选择
print("\n请选择一个问题：")
for key, value in preset_questions.items():
    print(f"{key}. {value}")

# 获取用户输入
try:
    user_choice = int(input("请输入问题编号："))
    if user_choice in preset_questions:
        user_query = preset_questions[user_choice]
        print(f"\n你选择了问题：{user_query}")
    else:
        print("无效的编号，将使用默认问题。")
        user_query = "用饼图分析员工职位分布"
except ValueError:
    print("输入无效，将使用默认问题。")
    user_query = "用饼图分析员工职位分布"

# 使用自然语言查询
result = smart_df.chat(user_query)
print(result)